'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-30 17:48:20
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-30 22:47:39
 * @Description: 注意力机制
'''

import paddle


class ChannelAttention(paddle.nn.Layer):

    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = paddle.nn.AdaptiveAvgPool2D(output_size=1)
        self.max_pool = paddle.nn.AdaptiveMaxPool2D(output_size=1)
        self.f1 = paddle.nn.Conv2D(in_channels=in_planes, out_channels=
            in_planes // ratio, kernel_size=1, bias_attr=False)
        self.relu = paddle.nn.ReLU()
        self.f2 = paddle.nn.Conv2D(in_channels=in_planes // ratio,
            out_channels=in_planes, kernel_size=1, bias_attr=False)
        self.act = paddle.nn.Sigmoid()

    def forward(self, x):
        avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
        max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
        out = self.act(avg_out + max_out)
        return out


class SpatialAttention(paddle.nn.Layer):

    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv = paddle.nn.Conv2D(in_channels=2, out_channels=1,
            kernel_size=kernel_size, padding=padding, bias_attr=False)
        self.act = paddle.nn.Sigmoid()

    def forward(self, x):
        avg_out = paddle.mean(x=x, axis=1, keepdim=True)
        max_out, _ = paddle.max(x=x, axis=1, keepdim=True), paddle.argmax(x
            =x, axis=1, keepdim=True)
        x = paddle.concat(x=[avg_out, max_out], axis=1)
        x = self.conv(x)
        return self.act(x)


class CBAMLayer(paddle.nn.Layer):

    def __init__(self, ch_in, ratio=16, kernel_size=7):
        super(CBAMLayer, self).__init__()
        self.channel_attention = ChannelAttention(ch_in, ratio)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        out = self.channel_attention(x) * x
        out = self.spatial_attention(out) * out
        return out


class SELayer(paddle.nn.Layer):

    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = paddle.nn.AdaptiveAvgPool2D(output_size=1)
        self.conv1 = paddle.nn.Conv2D(in_channels=channel, out_channels=
            channel // reduction, kernel_size=1, padding=0)
        self.relu = paddle.nn.ReLU()
        self.conv2 = paddle.nn.Conv2D(in_channels=channel // reduction,
            out_channels=channel, kernel_size=1, padding=0)
        self.sigmoid = paddle.nn.Sigmoid()

    def forward(self, x):
        module_input = x
        x = self.avg_pool(x)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.sigmoid(x)
        return module_input * x
